package cn.cloud.all.security.crypto.password;

import java.security.MessageDigest;
import java.util.Base64;

import cn.cloud.all.security.crypto.codec.Utf8;
import cn.cloud.all.security.crypto.keygen.BytesKeyGenerator;
import cn.cloud.all.security.crypto.keygen.KeyGenerators;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.bouncycastle.crypto.generators.SCrypt;

public class SCryptPasswordEncoder implements PasswordEncoder {

    private final Log logger = LogFactory.getLog(getClass());

    private final int cpuCost;

    private final int memoryCost;

    private final int parallelization;

    private final int keyLength;

    private final BytesKeyGenerator saltGenerator;

    public SCryptPasswordEncoder() {
        this(16384, 8, 1, 32, 64);
    }

    public SCryptPasswordEncoder(int cpuCost, int memoryCost, int parallelization, int keyLength, int saltLength) {
        if (cpuCost <= 1) {
            throw new IllegalArgumentException("Cpu cost parameter must be > 1.");
        }
        if (memoryCost == 1 && cpuCost > 65536) {
            throw new IllegalArgumentException("Cpu cost parameter must be > 1 and < 65536.");
        }
        if (memoryCost < 1) {
            throw new IllegalArgumentException("Memory cost must be >= 1.");
        }
        int maxParallel = Integer.MAX_VALUE / (128 * memoryCost * 8);
        if (parallelization < 1 || parallelization > maxParallel) {
            throw new IllegalArgumentException("Parallelisation parameter p must be >= 1 and <= " + maxParallel + " (based on block size r of " + memoryCost + ")");
        }
        if (keyLength < 1) {
            throw new IllegalArgumentException("Key length must be >= 1 and <= " + Integer.MAX_VALUE);
        }
        if (saltLength < 1) {
            throw new IllegalArgumentException("Salt length must be >= 1 and <= " + Integer.MAX_VALUE);
        }

        this.cpuCost = cpuCost;
        this.memoryCost = memoryCost;
        this.parallelization = parallelization;
        this.keyLength = keyLength;
        this.saltGenerator = KeyGenerators.secureRandom(saltLength);
    }

    public String encode(CharSequence rawPassword) {
        return digest(rawPassword, saltGenerator.generateKey());
    }

    public boolean matches(CharSequence rawPassword, String encodedPassword) {
        if (encodedPassword == null || encodedPassword.length() < keyLength) {
            logger.warn("Empty encoded password");
            return false;
        }
        return decodeAndCheckMatches(rawPassword, encodedPassword);
    }

    @Override
    public boolean upgradeEncoding(String encodedPassword) {
        if (encodedPassword == null || encodedPassword.isEmpty()) {
            return false;
        }

        String[] parts = encodedPassword.split("\\$");

        if (parts.length != 4) {
            throw new IllegalArgumentException("Encoded password does not look like SCrypt: " + encodedPassword);
        }

        long params = Long.parseLong(parts[1], 16);

        int cpuCost = (int) Math.pow(2, params >> 16 & 0xffff);
        int memoryCost = (int) params >> 8 & 0xff;
        int parallelization = (int) params & 0xff;

        return cpuCost < this.cpuCost
                || memoryCost < this.memoryCost
                || parallelization < this.parallelization;
    }

    private boolean decodeAndCheckMatches(CharSequence rawPassword, String encodedPassword) {
        String[] parts = encodedPassword.split("\\$");

        if (parts.length != 4) {
            return false;
        }

        long params = Long.parseLong(parts[1], 16);
        byte[] salt = decodePart(parts[2]);
        byte[] derived = decodePart(parts[3]);

        int cpuCost = (int) Math.pow(2, params >> 16 & 0xffff);
        int memoryCost = (int) params >> 8 & 0xff;
        int parallelization = (int) params & 0xff;

        byte[] generated = SCrypt.generate(Utf8.encode(rawPassword), salt, cpuCost, memoryCost, parallelization, keyLength);
        return MessageDigest.isEqual(derived, generated);
    }

    private String digest(CharSequence rawPassword, byte[] salt) {
        byte[] derived = SCrypt.generate(Utf8.encode(rawPassword), salt, cpuCost, memoryCost, parallelization, keyLength);
        String params = Long.toString(((int) (Math.log(cpuCost) / Math.log(2)) << 16L) | memoryCost << 8 | parallelization, 16);
        return "$" + params + '$' + encodePart(salt) + '$' + encodePart(derived);
    }

    private byte[] decodePart(String part) {
        return Base64.getDecoder().decode(Utf8.encode(part));
    }

    private String encodePart(byte[] part) {
        return Utf8.decode(Base64.getEncoder().encode(part));
    }
}
